Merged
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Contributor
Greptile OverviewGreptile SummaryThis PR adds a new Key changes:
Implementation notes: Confidence Score: 4/5
|
| Filename | Overview |
|---|---|
| litellm/proxy/_types.py | Added KEY_RESET_SPEND route enum and ResetSpendRequest model for the new endpoint. Changes are straightforward and follow existing patterns. |
| litellm/proxy/management_endpoints/key_management_endpoints.py | Implements reset_key_spend_fn endpoint with validation and authorization helpers. Properly validates input, checks permissions, updates database, and invalidates cache. Minor concern: doesn't update budget table's aggregate spend if shared budgets track total spend. |
| tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py | Comprehensive test coverage for reset_key_spend_fn including validation, authorization, success cases, edge cases, and error handling. Tests are well-structured and thorough. |
Sequence Diagram
sequenceDiagram
participant Client
participant Endpoint as reset_key_spend_fn
participant Validator as _validate_reset_spend_value
participant AuthCheck as _check_proxy_or_team_admin_for_key
participant DB as Prisma Client
participant Cache as user_api_key_cache
Client->>Endpoint: POST /key/{key}/reset_spend<br/>{reset_to: float}
Endpoint->>Endpoint: Check prisma_client exists
Endpoint->>Endpoint: Hash key if starts with "sk"
Endpoint->>DB: find_unique(token, include budget_table)
DB-->>Endpoint: Return key_in_db or None
alt Key not found
Endpoint-->>Client: 404 Key not found
end
Endpoint->>Validator: _validate_reset_spend_value(reset_to, key_in_db)
Validator->>Validator: Check type, non-negative, <= current_spend
Validator->>Validator: Check against max_budget (key & budget table)
alt Validation fails
Validator-->>Client: 400 Validation error
end
Validator-->>Endpoint: Return validated reset_to
Endpoint->>AuthCheck: _check_proxy_or_team_admin_for_key
AuthCheck->>AuthCheck: Check if PROXY_ADMIN
alt Not PROXY_ADMIN
AuthCheck->>DB: get_team_object(team_id)
DB-->>AuthCheck: Return team_table
AuthCheck->>AuthCheck: Check if user is team admin
end
alt Not authorized
AuthCheck-->>Client: 403 Forbidden
end
AuthCheck-->>Endpoint: Authorization passed
Endpoint->>DB: update(token, spend=reset_to)
DB-->>Endpoint: Return updated_key
Endpoint->>Cache: _delete_cache_key_object(hashed_token)
Cache-->>Endpoint: Cache invalidated
Endpoint-->>Client: 200 {spend, previous_spend, max_budget, etc.}
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Relevant issues
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unitCI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🆕 New Feature
✅ Test
Changes
Adds a new
/key/{key_id}/reset_spendendpoint that allows proxy admins and team admins to reset a key's spend to a specific value. The endpoint validates that the reset value is non-negative, does not exceed the current spend, and respects budget limits. Includes authorization checks to ensure only proxy admins or team admins can reset spend for keys in their teams.Screenshots